package com.c2b.algorithm.leetcode.base;

import java.util.*;

/**
 * <a href='https://leetcode.cn/problems/all-nodes-distance-k-in-binary-tree/'>二叉树中所有距离为 K 的结点(All Nodes Distance K in Binary Tree)</a>
 * <p>给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 k 。</p>
 * <p>返回到目标结点 target 距离为 k 的所有结点的值的列表。 答案可以以 任何顺序 返回。</p>
 *
 * <p>
 * <b>示例：</b>
 * <pre>
 * 示例 1：
 *      输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
 *                       3
 *                    /     \
 *                   <a color="red">5</a>       <a color="blue">1</a>
 *                  / \     / \
 *                 6   2   0   8
 *                    / \
 *                   <a color="blue">7</a>   <a color="blue">4</a>
 *      输出：[7,4,1]
 *      解释：所求结点为与目标结点（值为 5）距离为 2 的结点，值分别为 7，4，以及 1
 *
 * 示例 2:
 *      输入: root = [1], target = 1, k = 3
 *      输出: []
 * </pre>
 * </p>
 *
 * <p>
 * <b>提示:</b>
 * <ul>
 *     <li>节点数在 [1, 500] 范围内</li>
 *     <li>0 <= Node.val <= 500</li>
 *     <li>Node.val 中所有值 不同</li>
 *     <li>目标结点 target 是树上的结点。</li>
 *     <li>0 <= k <= 1000</li>
 * </ul>
 * </p>
 *
 * @author c2b
 * @since 2023/5/25 11:39
 */
public class LC0863AllNodesDistanceKInBinaryTree_M {


    static class Solution {
        public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
            List<Integer> allDistanceKList = new ArrayList<>();
            Map<TreeNode, TreeNode> parentNodeMap = new HashMap<>();
            buildParentNodeMap(root, parentNodeMap);
            find(target, k, new HashSet<>(), parentNodeMap, allDistanceKList);
            return allDistanceKList;
        }

        /**
         * 递归查询距离target节点为k的所有节点
         *
         * @param target           目标节点
         * @param k                距离k
         * @param visitedSet       已经访问过的节点
         * @param parentNodeMap    当前节点 与 父节点的映射关系
         * @param allDistanceKList 符合条件的节点值集合
         */
        private void find(TreeNode target, int k, HashSet<TreeNode> visitedSet, Map<TreeNode, TreeNode> parentNodeMap, List<Integer> allDistanceKList) {
            // 如果当前节点为 null，或者当前节点已经处理过。直接返回
            if (target == null || visitedSet.contains(target)) {
                return;
            }
            // 标识当前节点已经处理过，下次遇到无需处理
            visitedSet.add(target);
            // 如果距离为0，说明找到了合适的节点；无需继续查找，返回即可
            if (k == 0) {
                allDistanceKList.add(target.val);
                return;
            }
            // 向左子树查找，查找左子节点距离为k-1的节点
            find(target.left, k - 1, visitedSet, parentNodeMap, allDistanceKList);
            // 向右子树查找，查找右子节点距离为k-1的节点
            find(target.right, k - 1, visitedSet, parentNodeMap, allDistanceKList);
            // 向父节点查找，查找父节点距离为k-1的节点
            find(parentNodeMap.get(target), k - 1, visitedSet, parentNodeMap, allDistanceKList);
        }

        private void buildParentNodeMap(TreeNode currNode, Map<TreeNode, TreeNode> parentNodeMap) {
            if (currNode == null) {
                return;
            }
            if (currNode.left != null) {
                parentNodeMap.put(currNode.left, currNode);
                buildParentNodeMap(currNode.left, parentNodeMap);
            }
            if (currNode.right != null) {
                parentNodeMap.put(currNode.right, currNode);
                buildParentNodeMap(currNode.right, parentNodeMap);
            }
        }

        private Map<TreeNode, TreeNode> buildParentMapByIteration(TreeNode root) {
            Map<TreeNode, TreeNode> parentMap = new HashMap<>();
            Queue<TreeNode> queue = new LinkedList<>();
            queue.offer(root);
            parentMap.put(root, null);
            while (!queue.isEmpty()) {
                TreeNode currNode = queue.poll();
                if (currNode.left != null) {
                    queue.offer(currNode.left);
                    parentMap.put(currNode.left, currNode);
                }
                if (currNode.right != null) {
                    queue.offer(currNode.right);
                    parentMap.put(currNode.right, currNode);
                }
            }
            return parentMap;
        }
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode(3);
        root.left = new TreeNode(5);
        root.right = new TreeNode(1);
        root.left.left = new TreeNode(6);
        root.left.right = new TreeNode(2);
        root.right.left = new TreeNode(0);
        root.right.right = new TreeNode(8);
        root.left.right.left = new TreeNode(7);
        root.left.right.right = new TreeNode(4);
        Solution solution = new Solution();
        Printer.printListInteger(solution.distanceK(root, root.left, 2));
    }
}
